// Copyright (c) 2024 Go-Frame-Lite

package websocket

import (
	"context"
	"encoding/json"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
)

const (
	// 写入超时时间
	writeWait = 10 * time.Second

	// 读取超时时间
	pongWait = 60 * time.Second

	// 发送ping的周期，必须小于pongWait
	pingPeriod = (pongWait * 9) / 10

	// 最大消息大小
	maxMessageSize = 512
)

var (
	newline = []byte("\n")
	space   = []byte(" ")
)

// Upgrader 配置websocket升级器
var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	// 允许所有CORS请求
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// Client 是websocket连接的中间人
// 每个Client代表一个websocket连接

type Client struct {
	// 管理器实例
	Hub *Hub

	// websocket连接
	conn *websocket.Conn

	// 发送消息的缓冲通道
	send chan []byte

	// 客户端ID
	ID string

	// 客户端元数据
	Metadata map[string]interface{}

	// 用于保护Metadata的互斥锁
	mutex sync.RWMutex
}

// Hub 维护活动客户端集合并广播消息

type Hub struct {
	// 已注册的客户端
	Clients map[*Client]bool

	// 客户端ID映射，方便通过ID查找客户端
	ClientsByID map[string]*Client

	// 来自客户端的入站消息
	broadcast chan []byte

	// 注册请求
	register chan *Client

	// 注销请求
	unregister chan *Client

	// 用于保护Clients和ClientsByID的互斥锁
	mutex sync.RWMutex
}

// Message 定义WebSocket消息结构
type Message struct {
	Type      string          `json:"type"`              // 消息类型
	Content   json.RawMessage `json:"content,omitempty"` // 消息内容
	SenderID  string          `json:"sender_id,omitempty"` // 发送者ID
	Timestamp int64           `json:"timestamp"`         // 时间戳
}

// UnmarshalJSON 实现json.Unmarshaler接口
func (m *Message) UnmarshalJSON(data []byte) error {
	type Alias Message
	aux := &struct {
		*Alias
	}{Alias: (*Alias)(m)}
	return json.Unmarshal(data, aux)
}

// HandlerFunc 定义WebSocket处理函数类型
type HandlerFunc func(*Client, []byte)

// NewHub 创建一个新的Hub实例
func NewHub() *Hub {
	return &Hub{
		broadcast:   make(chan []byte),
		register:    make(chan *Client),
		unregister:  make(chan *Client),
		Clients:     make(map[*Client]bool),
		ClientsByID: make(map[string]*Client),
	}
}

// Run 启动Hub的主循环
func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.mutex.Lock()
			h.Clients[client] = true
			h.ClientsByID[client.ID] = client
			h.mutex.Unlock()

		case client := <-h.unregister:
			h.mutex.Lock()
			if _, ok := h.Clients[client]; ok {
				delete(h.Clients, client)
				delete(h.ClientsByID, client.ID)
				close(client.send)
			}
			h.mutex.Unlock()

		case message := <-h.broadcast:
			h.mutex.RLock()
			for client := range h.Clients {
				select {
				case client.send <- message:
				default:
					close(client.send)
					delete(h.Clients, client)
					delete(h.ClientsByID, client.ID)
				}
			}
			h.mutex.RUnlock()
		}
	}
}

// Broadcast 广播消息给所有客户端
func (h *Hub) Broadcast(message []byte) {
	h.broadcast <- message
}

// BroadcastJSON 广播JSON格式的消息
func (h *Hub) BroadcastJSON(msgType string, content interface{}) error {
	// 将content转换为json.RawMessage
	jsonContent, err := json.Marshal(content)
	if err != nil {
		return err
	}
	
	msg := Message{
		Type:      msgType,
		Content:   jsonContent,
		Timestamp: time.Now().UnixNano() / int64(time.Millisecond),
	}

	data, err := json.Marshal(msg)
	if err != nil {
		return err
	}

	h.Broadcast(data)
	return nil
}

// GetClient 通过ID获取客户端
func (h *Hub) GetClient(id string) *Client {
	h.mutex.RLock()
	defer h.mutex.RUnlock()
	return h.ClientsByID[id]
}

// GetClientsCount 获取当前连接的客户端数量
func (h *Hub) GetClientsCount() int {
	h.mutex.RLock()
	defer h.mutex.RUnlock()
	return len(h.Clients)
}

// 读取客户端消息
func (c *Client) readPump() {
	defer func() {
		c.Hub.unregister <- c
		c.conn.Close()
	}()

	c.conn.SetReadLimit(maxMessageSize)
	c.conn.SetReadDeadline(time.Now().Add(pongWait))
	c.conn.SetPongHandler(func(string) error {
		c.conn.SetReadDeadline(time.Now().Add(pongWait))
		return nil
	})

	for {
		_, message, err := c.conn.ReadMessage()
		if err != nil {
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("error: %v", err)
			}
			break
		}

		// 处理接收到的消息
		handleMessage(c, message)
	}
}

// 发送消息给客户端
func (c *Client) writePump() {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.conn.Close()
	}()

	for {
		select {
		case message, ok := <-c.send:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if !ok {
				// Hub关闭了通道
				c.conn.WriteMessage(websocket.CloseMessage, []byte{})
				return
			}

			w, err := c.conn.NextWriter(websocket.TextMessage)
			if err != nil {
				return
			}
			w.Write(message)

			// 添加队列中的所有消息
			n := len(c.send)
			for i := 0; i < n; i++ {
				w.Write(newline)
				w.Write(<-c.send)
			}

			if err := w.Close(); err != nil {
				return
			}
		case <-ticker.C:
			c.conn.SetWriteDeadline(time.Now().Add(writeWait))
			if err := c.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

// Send 发送消息给客户端
func (c *Client) Send(message []byte) {
	c.send <- message
}

// SendJSON 发送JSON格式消息给客户端
func (c *Client) SendJSON(msgType string, content interface{}) error {
	// 将content转换为json.RawMessage
	jsonContent, err := json.Marshal(content)
	if err != nil {
		return err
	}
	
	msg := Message{
		Type:      msgType,
		Content:   jsonContent,
		Timestamp: time.Now().UnixNano() / int64(time.Millisecond),
	}

	data, err := json.Marshal(msg)
	if err != nil {
		return err
	}

	c.Send(data)
	return nil
}

// SetMetadata 设置客户端元数据
func (c *Client) SetMetadata(key string, value interface{}) {
	c.mutex.Lock()
	defer c.mutex.Unlock()
	if c.Metadata == nil {
		c.Metadata = make(map[string]interface{})
	}
	c.Metadata[key] = value
}

// GetMetadata 获取客户端元数据
func (c *Client) GetMetadata(key string) (interface{}, bool) {
	c.mutex.RLock()
	defer c.mutex.RUnlock()
	if c.Metadata == nil {
		return nil, false
	}
	value, exists := c.Metadata[key]
	return value, exists
}

// 消息处理函数
var messageHandlers = make(map[string]HandlerFunc)

// RegisterHandler 注册消息处理器
func RegisterHandler(msgType string, handler HandlerFunc) {
	messageHandlers[msgType] = handler
}

// handleMessage 处理接收到的消息
func handleMessage(client *Client, data []byte) {
	var msg Message
	if err := json.Unmarshal(data, &msg); err != nil {
		log.Printf("解析消息失败: %v", err)
		return
	}

	// 设置发送者ID
	msg.SenderID = client.ID

	// 查找对应的处理器
	if handler, ok := messageHandlers[msg.Type]; ok {
		handler(client, data)
	} else {
		log.Printf("未注册的消息类型: %s", msg.Type)
	}
}

// Handler 是WebSocket路由处理函数
func Handler(c *gin.Context, hub *Hub) {
	// 升级HTTP连接为WebSocket
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Println(err)
		return
	}

	// 创建客户端实例
	clientID := c.Query("client_id")
	if clientID == "" {
		clientID = c.Request.RemoteAddr
	}

	client := &Client{
		Hub:      hub,
		conn:     conn,
		send:     make(chan []byte, 256),
		ID:       clientID,
		Metadata: make(map[string]interface{}),
	}

	// 注册客户端
	client.Hub.register <- client

	// 启动goroutine处理读写
	go client.writePump()
	go client.readPump()

	// 发送欢迎消息
	welcomeMsg := map[string]string{
		"message": "连接成功",
		"client_id": client.ID,
	}
	client.SendJSON("welcome", welcomeMsg)
}

// GinHandler 返回一个Gin处理函数
func GinHandler(hub *Hub) gin.HandlerFunc {
	return func(c *gin.Context) {
		Handler(c, hub)
	}
}

// WithContext 在上下文中添加WebSocket Hub
func WithContext(ctx context.Context, hub *Hub) context.Context {
	return context.WithValue(ctx, "websocket_hub", hub)
}

// FromContext 从上下文中获取WebSocket Hub
func FromContext(ctx context.Context) (*Hub, bool) {
	hub, ok := ctx.Value("websocket_hub").(*Hub)
	return hub, ok
}